# _*_ coding: utf-8 _*_
"""
Time:     2024/9/14
Author:   2051261
File:     operate_columns.py
"""
from __future__ import annotations
import typing as t

from schematics.models import Model
from schematics.types import StringType
from schematics.types import ListType
from schematics.types import DictType
from schematics.types import BaseType
from schematics.types import ModelType

from extractor.const import ContextKey
from extractor.flow.node_model.child_node_method.check import validate_data
from extractor.flow.node_model.child_node_method.modify import modify_element
from extractor.flow.base_node_flow import BaseNodeFlow

if t.TYPE_CHECKING:
    from extractor.flow.node_model.log_node import RecordLog
    from extractor.flow.task_entry import Cache

operator_columns = {
    "check_element": validate_data,
    "modify_element": modify_element,
}


class ChildModel(Model):
    operate_func = StringType(required=True)
    kwargs = DictType(BaseType(required=False), required=True)  # 允许参数值为None


class OperateColumnModel(Model):
    check_element = ListType(ModelType(ChildModel))
    modify_element = ListType(ModelType(ChildModel))


class OperateColumnNode(BaseNodeFlow):

    async def __call__(self, record_log: RecordLog, context: t.Iterator[t.Dict[str, t.Any]], cache: Cache,
                       record: t.DefaultDict, kwargs=None) -> t.Iterator[t.Dict[str, t.Any]]:

        for item in context:
            data = item[ContextKey.DATA]
            tmp = {ContextKey.DATA: data}
            for operate_type, operate_list, in kwargs.items():
                result = await operator_columns[operate_type](data, operate_list)
                tmp.update(result)
                if tmp.get(ContextKey.DIRTY):
                    break
            yield tmp
